import os
import re
from pathlib import Path
from openpyxl import Workbook
from openpyxl.styles import Font, Border, Side
from openpyxl.utils import get_column_letter

class ExtensionManager:
    def __init__(self, base_dir="extensions"):
        self.base_dir = Path(base_dir).resolve()
        self.extensions = {}  # {扩展名称: [指令1, 指令2, ...]}
        self._load_extensions()

    def _load_extensions(self):
        """加载所有扩展文件及指令，过滤注释行"""
        try:
            if not self.base_dir.exists():
                print(f"警告: 扩展目录 '{self.base_dir}' 不存在，已创建空目录")
                self.base_dir.mkdir(parents=True, exist_ok=True)
                return

            for file_path in self.base_dir.rglob("*"):
                if file_path.is_file():
                    rel_path = file_path.relative_to(self.base_dir)
                    ext_name = str(rel_path).replace(os.sep, "_")

                    try:
                        with open(file_path, "r", encoding="utf-8") as f:
                            # 过滤空行和#开头的注释行
                            instructions = [
                                line.strip()
                                for line in f
                                if line.strip() and not line.strip().startswith('#')
                            ]

                        self.extensions[ext_name] = instructions
                        print(f"加载扩展: {ext_name} (含 {len(instructions)} 条指令)")

                    except UnicodeDecodeError:
                        print(f"警告: 无法解码文件 '{file_path}'，已跳过")
                    except Exception as e:
                        print(f"处理文件 '{file_path}' 时出错: {str(e)}")

        except PermissionError:
            print(f"错误: 没有访问目录 '{self.base_dir}' 的权限")
        except Exception as e:
            print(f"加载扩展时出错: {str(e)}")

    def list_extensions(self):
        """打印所有扩展名称（数组名）"""
        with open('dist/vt_ext', 'w', encoding='utf-8') as file:
          file.write("\n===== 所有指令扩展 =====")
          if not self.extensions:
              file.write("\n没有加载到任何扩展指令")
              return

          for i, ext_name in enumerate(sorted(self.extensions.keys()), 1):
              file.write(f"\n{i}. {ext_name} (指令数: {len(self.extensions[ext_name])})")

    def extract_field(self, line):
        """
        提取指令字段，按规则处理：
        1. 过滤#开头的注释行（已在加载时处理）
        2. 遇到$import提取第2个字段（索引1）
        3. 遇到$pseudo_op提取第3个字段（索引2）
        4. 默认返回第一个字段（索引0）
        """
        line = line.strip()
        if not line:
            return None

        parts = line.split()
        if not parts:
            return None

        info = []
        # 处理$import指令（提取第2个字段，索引1）
        if parts[0] == '$import':
            info = [parts[1], "ref", "ref"]
        # 处理$pseudo_op指令（提取第3个字段，索引2）
        elif parts[0] == '$pseudo_op':
            if len(parts) >= 3:  # 至少需要3个字段
                info = parts[2:5]
            else:
                info = [parts[2], "ref", "ref"]
        else:
            info = parts[0:3]

        # 默认返回第一个字段
        return info

    def process_instruction_fields(self, ext_name, debug=0):
        """处理指令并按规则提取字段，同时展示完整指令"""
        if ext_name not in self.extensions:
            print(f"错误: 扩展 '{ext_name}' 不存在")
            return []

        instructions = self.extensions[ext_name]
        result = []
        if debug:
            print(f"\n===== 扩展 '{ext_name}' 的字段提取结果 =====")

        if not instructions:
            print("该扩展没有有效指令")
            return result

        for i, instr in enumerate(instructions, 1):
            fields = self.extract_field(instr)
            result.append((fields, instr))
            if debug:
                print(f"{i}. 指令: {fields} -> 完整指令: {instr}")

        return result

    def export_to_excel(self, filename="riscv_extensions_fields.xlsx"):
        """导出所有扩展的提取字段结果到Excel"""
        if not self.extensions:
            print("没有可导出的扩展数据")
            return

        wb = Workbook()
        ws = wb.active
        ws.title = "RISC-V指令扩展字段提取"

        # 样式设置
        header_font = Font(bold=True, size=12)
        ext_font = Font(bold=True, color="00008B")
        thin_border = Border(
            left=Side(style='thin'), right=Side(style='thin'),
            top=Side(style='thin'), bottom=Side(style='thin')
        )

        # 表头
        headers = ["扩展名称", "指令", "初级执行单元", "次级执行单元", "完整指令"]
        ws.append(headers)
        for col_idx, header in enumerate(headers, 1):
            cell = ws.cell(row=1, column=col_idx)
            cell.font = header_font
            cell.border = thin_border

        # 填充数据
        current_row = 2
        for ext_name in sorted(self.extensions.keys()):
            instructions = self.extensions[ext_name]
            if not instructions:
                continue

            # 处理扩展名称合并单元格
            instr_count = len(instructions)
            start_row = current_row
            end_row = current_row + instr_count - 1

            if instr_count > 1:
                ws.merge_cells(
                    start_row=start_row,
                    start_column=1,
                    end_row=end_row,
                    end_column=1
                )

            # 仅对左上角主单元格赋值
            ext_cell = ws.cell(row=start_row, column=1)
            ext_cell.value = ext_name
            ext_cell.font = ext_font
            ext_cell.border = thin_border

            # 处理指令字段
            processed = self.process_instruction_fields(ext_name, 1)
            for idx, (fields, instr) in enumerate(processed):
                row = start_row + idx

                # 指令
                field0_cell = ws.cell(row=row, column=2, value=fields[0])
                field0_cell.border = thin_border

                field1_cell = ws.cell(row=row, column=3, value=fields[1])
                field1_cell.border = thin_border

                field2_cell = ws.cell(row=row, column=4, value=fields[2])
                field2_cell.border = thin_border

                # 完整指令
                instr_cell = ws.cell(row=row, column=5, value=instr)
                instr_cell.border = thin_border

            current_row = end_row + 1

        # 自动调整列宽
        for col_idx in range(1, 4):
            max_length = 0
            column_letter = get_column_letter(col_idx)

            for cell in ws[column_letter]:
                try:
                    if len(str(cell.value)) > max_length:
                        max_length = len(cell.value)
                except:
                    pass

            adjusted_width = (max_length + 2) * 1.2
            ws.column_dimensions[column_letter].width = min(adjusted_width, 100)

        try:
            wb.save(filename)
            print(f"\n数据已成功导出到 {os.path.abspath(filename)}")
        except PermissionError:
            print(f"错误: 没有权限写入文件 '{filename}'，可能文件已被打开")
        except Exception as e:
            print(f"导出Excel时出错: {str(e)}")

    def export_insns_to_txt(self):
        """导出所有指令数据到txt"""
        if not self.extensions:
            print("没有可导出的扩展数据")
            return

        with open('dist/insns_info.txt', 'w', encoding='utf-8') as file:
            for ext_name in sorted(self.extensions.keys()):
                instructions = self.extensions[ext_name]
                if not instructions:
                    continue

                prefix_pattern = re.compile(r'^(rv_|rv64_|rv32_)')
                rv_ext_name = prefix_pattern.sub('', ext_name)
                # 处理指令字段
                processed = self.process_instruction_fields(ext_name)
                for idx, (fields, instr) in enumerate(processed):
                    file.write(f'{fields[0]} {fields[1]} {fields[2]} {rv_ext_name}\n')

        print(f'\n数据已成功导出到 {os.path.abspath("dist/insns_info.txt")}')

def main():
    ext_manager = ExtensionManager()

    ext_manager.list_extensions()

    filename = "dist/insn_execution_unit_map.xlsx"
    ext_manager.export_to_excel(filename)

    ext_manager.export_insns_to_txt()

    print("\n===== RISC-V指令扩展处理工具 =====")
    print("1. 列出所有扩展名称")
    print("2. 导出所有扩展的指令结果到Excel")
    print("3. 导出所有扩展的指令结果到txt")

if __name__ == "__main__":
    main()
